
# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from client import sys, time, os, gc,n,ws,wx,math
from client import matplotlib,rcParams,cm,Toolbar,FigureCanvasWxAgg,NavigationToolbar2WxAgg
from client import Figure,nx,RectangleSelector,_load_bitmap,rand

import matplotlib.widgets as mw
import matplotlib.numerix as numerix
from matplotlib.numerix import arange, sin, cos, pi
import random

TIMER_ID = wx.NewId()

class MyNavigationToolbar(NavigationToolbar2WxAgg):
	"""
	Extend the default wx toolbar with your own event handlers
	"""
	ON_UP = wx.NewId()
	ON_DOWN = wx.NewId()
	ON_ALL = wx.NewId()
	def __init__(self, parent, canvas, cankill):
		NavigationToolbar2WxAgg.__init__(self, canvas)
		self.parent = parent
		if self.parent.spec_type == 'iv':
			self.AddSimpleTool(self.ON_DOWN, _load_bitmap('stock_down.xpm'),
						'Move down a z-layer', 'Move down a z-layer')
			wx.EVT_TOOL(self, self.ON_DOWN, self._on_down)
			self.AddSimpleTool(self.ON_UP, _load_bitmap('stock_up.xpm'),
						'Move up a z-layer', 'Move up a z-layer')
			wx.EVT_TOOL(self, self.ON_UP, self._on_up)
			self.AddSimpleTool(self.ON_ALL, _load_bitmap('stock_refresh.xpm'),'Display All', 'Display All')
			wx.EVT_TOOL(self, self.ON_ALL, self._on_all)
		else: pass

	def _on_up(self, evt):
		self.parent.index += 1
		if self.parent.index > self.parent.max_index: self.parent.index = self.parent.max_index
		else: self.parent.plotSelection(self.parent.index)
		self.parent.mode=0
		evt.Skip()
	
	def _on_down(self, evt):
		self.parent.index -= 1
		if self.parent.index < 0: self.parent.index = 0
		else: self.parent.plotSelection(self.parent.index)
		self.parent.mode=1
		evt.Skip()
		
	def _on_all(self, evt):
		if self.parent.pdir == 0: self.parent.pdir = 1
		else: self.parent.pdir = 0
		self.parent.plotAll(self.parent.pdir)
		self.parent.mode=2
		evt.Skip()

class SpecFigure(wx.Frame,MyNavigationToolbar):
	ASPECT_W = ws.ASPECT_W
	ASPECT_H = ws.ASPECT_H
	ASPECT_SIZE = ws.ASPECT_SIZE
	TITLE_SIZE = ws.TITLE_SIZE
	AXES_SIZE = ws.AXES_SIZE
	DEFAULT_COLOR = ws.DEFAULT_COLOR
	POINT_TYPE = ws.POINT_TYPE
	LINE_WEIGHT = ws.LINE_WEIGHT
	pdir = 0
	mode = 0
	def __init__(self,parent,name,spec_type='iv',external=None,toolbar=None):
		self.parent = parent
		if not external:
			wx.Frame.__init__(self, parent, -1, name, style=wx.STAY_ON_TOP|wx.DEFAULT_FRAME_STYLE)
			self.fig = Figure((self.ASPECT_W,self.ASPECT_H), self.ASPECT_SIZE)
			self.canvas = FigureCanvasWxAgg(self, -1, self.fig)
			self.t = wx.Timer(self, TIMER_ID)
			self.t.Start(10000)
			self.txt = None
			self.external = False
		else:
			#self = external[0]
			self.fig = external[1]
			self.canvas = external[2]
			self.external = True
		self.spec_type = spec_type
		#self.fig = Figure(None, None)
		#self.canvas = FigureCanvasWxAgg(self, -1, self.fig)
		#if not canvas: self.canvas = NoRepaintCanvas(self, -1, self.fig)
		#else: self.canvas = canvas
		#self.Bind(wx.EVT_IDLE, self._onIdle)
		#self.Bind(wx.EVT_SIZE, self._onSize)
		self._resizeflag = True
		if toolbar == None:
			self.toolbar = MyNavigationToolbar(self, self.canvas, True)
			self.toolbar.Realize()
			tw, th = self.toolbar.GetSizeTuple()
			fw, fh = self.canvas.GetSizeTuple()
			self.toolbar.SetSize(wx.Size(fw, th))
		else:
			self.toolbar = toolbar
		self.x = None
		self.y = None
		self.active = True
		self.index = 0
		self.max_index = 0
		
		# On Windows, default frame size behaviour is incorrect
		# you don't need this under Linux
		
		self.init_plot_data()
		# Initialize the timer - wxPython requires this to be connected to
		# the receiving event handler
		
		# Now put all into a sizer
		sizer = wx.BoxSizer(wx.VERTICAL)
		
		#sizer = wx.BoxSizer(wx.HORIZONTAL)
		#self.SetSizer(sizer)
		#sizer.SetItemMinSize(self, 300, 300)
		#self.Fit()
		if not external:
			# This way of adding to sizer allows resizing
			sizer.Add(self.canvas, 1, wx.LEFT|wx.TOP|wx.GROW)
			# Best to allow the toolbar to resize!
			sizer.Add(self.toolbar, 0, wx.GROW)
			self.SetSizer(sizer)
			self.Fit()
			self.canvas.Bind(wx.EVT_RIGHT_DOWN, self.OnRightDown)
			self.Bind(wx.EVT_TIMER, self.onTimer, id=TIMER_ID)
			self.Bind(wx.EVT_CLOSE, self.onClose)
			
			# Note that event is a MplEvent
			self.canvas.mpl_connect('motion_notify_event', self.UpdateStatusBar)
			self.canvas.Bind(wx.EVT_ENTER_WINDOW, self.ChangeCursor)
			
			self.statusBar = wx.StatusBar(self, -1)
			self.statusBar.SetFieldsCount(1)
			self.SetStatusBar(self.statusBar)
			
			#self._SetSize()
			self.Show()

	def init_plot_data(self):
		# jdh you can add a subplot directly from the fig rather than
		# the fig manager
		self.fig.clear()
		if not hasattr(self, 'subplot'):
                	self.subplot = self.fig.add_subplot(111)
		self.fig.subplots_adjust(left=0.15, bottom=0.2)
		theta = arange(0, 45*2*pi, 0.02)
		rad = (0.8*theta/(2*pi)+1)
		r = rad*(8 + sin(theta*7+rad/1.8))
		self.Zx = r*cos(theta)
		self.Zy = r*sin(theta)
		#Now draw it
		self.subplot.plot(self.Zx,self.Zy, self.POINT_TYPE)
		#Set some plot attributes
		self.subplot.set_title("A Typical IV Curve %s"%len(self.Zx), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Sample Bias (V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = self.AXES_SIZE)
		self.subplot.set_xlim([-400, 400])
		self.subplot.set_ylim([-400, 400])
	
	def _drawOnChange(self,coords,data):
		self.data = data
		self.coords = coords
		if self.spec_type == 'iv':
			if self.mode == 0: self._drawIVOnChange(coords,data)
			elif self.mode == 1: self.plotSelection(self.index)
			elif self.mode == 2: self.plotAll(self.pdir)
		elif self.spec_type == 'zv':
			self._drawZVOnChange(coords,data)
		elif self.spec_type == 'iz':
			self._drawIZOnChange(coords,data)
		else: raise 'invalid spectrosopy type'
	
	def _drawIVOnChange(self,coords,data):
		self.data = data
		self.max_index = len(self.data)-1
		self.coords = coords
		colors = ['g','r','b','c','y','m']
		xmax,xmin,ymax,ymin = (-10,10,-10,10)
		for j in range(2):
			if xmax < n.array(data[0][j][0]).max(): xmax = n.array(data[0][j][0]).max()
			if xmin > n.array(data[0][j][0]).min(): xmin = n.array(data[0][j][0]).min()
			if ymax < n.array(data[0][j][1]).max(): ymax = n.array(data[0][j][1]).max()
			if ymin > n.array(data[0][j][1]).min(): ymin = n.array(data[0][j][1]).min()
		if abs(ymax - ymin)<0.1:
			ymax += 0.01
			ymin -= 0.01
		self.subplot.clear()
		z1f = self.subplot.plot(data[0][0][0],data[0][0][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[0],lw=self.LINE_WEIGHT)
		z1r = self.subplot.plot(data[0][1][0],data[0][1][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[1],lw=self.LINE_WEIGHT)
		if not self.external: self.fig.subplots_adjust(left=0.15, bottom=0.2)
		else: self.fig.subplots_adjust(left=0.1, bottom=0.1,top=0.85,right=0.9)
		axcolor = 'w'
		#Set some plot attributes

		self.subplot.set_title("IV Curve at %s"%str(coords), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Sample Bias (V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = self.AXES_SIZE)
		self.subplot.set_xlim([xmin, xmax])
		self.subplot.set_ylim([ymin, ymax])
		self.subplot.legend([z1f,z1r],['z0f', 'z0r'], shadow=True)
		self.canvas.draw()

	def _drawZVOnChange(self,coords,data):
		self.data = data
		self.coords = coords
		colors = ['g','r','b','c','y','m']
		xmax,xmin,ymax,ymin = (-10,10,-10,10)
		for j in range(2):
			if xmax < n.array(data[0][j][0]).max(): xmax = n.array(data[0][j][0]).max()
			if xmin > n.array(data[0][j][0]).min(): xmin = n.array(data[0][j][0]).min()
			if ymax < n.array(data[0][j][1]).max(): ymax = n.array(data[0][j][1]).max()
			if ymin > n.array(data[0][j][1]).min(): ymin = n.array(data[0][j][1]).min()
		if abs(ymax - ymin)<0.1:
			ymax += 0.01
			ymin -= 0.01
		self.subplot.clear()
		z1f = self.subplot.plot(data[0][0][0],data[0][0][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[0],lw=self.LINE_WEIGHT)
		z1r = self.subplot.plot(data[0][1][0],data[0][1][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[1],lw=self.LINE_WEIGHT)
		if not self.external: self.fig.subplots_adjust(left=0.15, bottom=0.2)
		else: self.fig.subplots_adjust(left=0.1, bottom=0.1,top=0.85,right=0.9)
		axcolor = 'w'
		#Set some plot attributes

		self.subplot.set_title("ZV Curve at %s"%str(coords), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Sample Bias (V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Z Position", fontsize = self.AXES_SIZE)
		self.subplot.set_xlim([xmin, xmax])
		self.subplot.set_ylim([ymin, ymax])
		self.subplot.legend([z1f,z1r],['forward', 'reverse'], shadow=True)
		self.canvas.draw()

	def _drawIZOnChange(self,coords,data):
		self.data = data
		self.coords = coords
		colors = ['g','r','b','c','y','m']
		xmax,xmin,ymax,ymin = (-10,10,-10,10)
		for j in range(2):
			if xmax < n.array(data[0][j][0]).max(): xmax = n.array(data[0][j][0]).max()
			if xmin > n.array(data[0][j][0]).min(): xmin = n.array(data[0][j][0]).min()
			if ymax < n.array(data[0][j][1]).max(): ymax = n.array(data[0][j][1]).max()
			if ymin > n.array(data[0][j][1]).min(): ymin = n.array(data[0][j][1]).min()
		if abs(ymax - ymin)<0.1:
			ymax += 0.01
			ymin -= 0.01
				
		self.subplot.clear()
		z1f = self.subplot.plot(data[0][0][0],data[0][0][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[0],lw=self.LINE_WEIGHT)
		z1r = self.subplot.plot(data[0][1][0],data[0][1][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[1],lw=self.LINE_WEIGHT)
		if not self.external: self.fig.subplots_adjust(left=0.15, bottom=0.2)
		else: self.fig.subplots_adjust(left=0.1, bottom=0.1,top=0.85,right=0.9)
		#axcolor = 'w'
		#Set some plot attributes
		self.subplot.set_title("IZ Curve %s"%str(coords), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Z Position (dZ, in V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = self.AXES_SIZE)

		self.subplot.set_xlim([xmin, xmax])
		self.subplot.set_ylim([ymin, ymax])
		self.subplot.legend([z1f,z1r],['forward', 'reverse'], shadow=True)
		self.canvas.draw()
		

	#def colorfunc(self,label):
		#index = int(label.strip('z'))
		#if not self.index_check == index:
			#self.index_check = index
			#print 'bleeehhhhhhh'
			#self.plotSelection(index)
		#print dir(event),event.index()
		#print event.x
		#l.set_color(label)
		#draw()
	
	def plotSelection(self,index):
		data = self.data
		colors = ['g','r','b','c','y','m']
		coords=self.coords
		xmax,xmin,ymax,ymin = (-10,10,-10,10)
		for j in range(2):
			if xmax < n.array(data[index][j][0]).max(): xmax = n.array(data[index][j][0]).max()
			if xmin > n.array(data[index][j][0]).min(): xmin = n.array(data[index][j][0]).min()
			if ymax < n.array(data[index][j][1]).max(): ymax = n.array(data[index][j][1]).max()
			if ymin > n.array(data[index][j][1]).min(): ymin = n.array(data[index][j][1]).min()
		if abs(ymax - ymin)<0.1:
			ymax += 0.01
			ymin -= 0.01
		self.subplot.clear()
		znf = self.subplot.plot(data[index][0][0],data[index][0][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[0],lw=self.LINE_WEIGHT)
		znr = self.subplot.plot(data[index][1][0],data[index][1][1], self.POINT_TYPE,color=self.DEFAULT_COLOR[1],lw=self.LINE_WEIGHT)
		if not self.external: self.fig.subplots_adjust(left=0.15, bottom=0.2)
		else: self.fig.subplots_adjust(left=0.1, bottom=0.1,top=0.85,right=0.9)
		axcolor = 'w'
		#Set some plot attributes
		self.subplot.set_title("IV Curve %s"%str(coords), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Sample Bias (V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = self.AXES_SIZE)
		#self.subplot.set_xlabel("Tip Bias (V)", fontsize = 16)
		#self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = 16)
		self.subplot.set_xlim([xmin, xmax])
		self.subplot.set_ylim([ymin, ymax])
		self.subplot.legend([znf,znr],['z'+str(index)+'f', 'z'+str(index)+'r'], shadow=True)
		self.canvas.draw()
	
	def plotAll(self,direction):
		data = self.data
		layers = len(data)
		colors = ['g','r','b','c','y','m']
		coords=self.coords
		xmax,xmin,ymax,ymin = (-10,10,-10,10)
		for i in range(layers):
			if xmax < n.array(data[i][direction][0]).max(): xmax = n.array(data[i][direction][0]).max()
			if xmin > n.array(data[i][direction][0]).min(): xmin = n.array(data[i][direction][0]).min()
			if ymax < n.array(data[i][direction][1]).max(): ymax = n.array(data[i][direction][1]).max()
			if ymin > n.array(data[i][direction][1]).min(): ymin = n.array(data[i][direction][1]).min()
		if abs(ymax - ymin)<0.1:
			ymax += 0.01
			ymin -= 0.01	
		self.subplot.clear()
		count = 0
		zn = []
		zind = []
		for i in range(layers):
			if direction == 0: direc = 'f'
			else: direc = 'r'
			zind.append('z'+str(i)+direc)
			zn.append(self.subplot.plot(data[i][direction][0],data[i][direction][1],\
			self.POINT_TYPE,color=self.DEFAULT_COLOR[count],lw=self.LINE_WEIGHT))
			count += 1
		if not self.external: self.fig.subplots_adjust(left=0.15, bottom=0.2)
		else: self.fig.subplots_adjust(left=0.1, bottom=0.1,top=0.85,right=0.9)
		axcolor = 'w'
		#Set some plot attributes
		self.subplot.set_title("IV Curve at %s"%str(coords), fontsize = self.TITLE_SIZE)
		self.subplot.set_xlabel("Sample Bias (V)", fontsize = self.AXES_SIZE)
		self.subplot.set_ylabel("Tunneling Current (mA)", fontsize = self.AXES_SIZE)
		self.subplot.set_xlim([xmin, xmax])
		self.subplot.set_ylim([ymin, ymax])
		self.subplot.legend(zn,zind, shadow=True)
		self.canvas.draw()
	
	#def _onSize(self, event):
		#self._resizeflag = True

	#def _onIdle(self, evt):
		#if self._resizeflag:
			#self._resizeflag = False
			#self._SetSize()
			##self.draw()

	#def _SetSize(self, pixels = None):
		#"""
		#This method can be called to force the Plot to be a desired size, which defaults to
		#the ClientSize of the panel
		#"""
		#if not pixels:
			#pixels = self.GetClientSize()
		#self.canvas.SetSize(pixels)
		#self.fig.set_size_inches(pixels[0]/self.fig.get_dpi(),
		#pixels[1]/self.fig.get_dpi())

	def ChangeCursor(self, event):
		self.canvas.SetCursor(wx.StockCursor(wx.CURSOR_ARROW))

	def UpdateStatusBar(self, event):
		if event.inaxes:
			self.x, self.y = event.xdata, event.ydata
			self.statusBar.SetStatusText(( "x= " + str(self.x)[:5] +
							"  y=" +str(self.y)[:5] ),
							0)

	def onTimer(self, evt):
		self.canvas.draw()

	def onClose(self, evt):
		self.active = False
		self.t.Stop()
		self.parent.specsub = None
		evt.Skip()

	def OnRightDown(self, evt):
		"""right mouse button is pressed"""
		#ax = self.canvas.figure.axes[0]
		#rgb = tuple(rand(3))
		#pt = evt.GetPosition()  # position tuple
		#self.txt = ax.text(0, 0, str(pt),
			#transform=ax.transAxes,
			#color=rgb)
		#self.canvas.draw()
		pass